
#include <vector>
#include "AESSymmetricKey.h"
#include "ECDSAPublicKey.h"
#include "DSAPublicKey.h"
#include "RSAPublicKey.h"
#include "RSAPrivateKey.h"
#include "CryptOChain.h"
#include "DecryptWrapper.h"

/** General initialisation function for the decryption wrapper.
 * 
 * @param deckey DecryptionKey to use for decrypting data if present
 * and not NULL.
 * 
 * @param chkkey CheckingKey to use to verify signatures on the data
 * stream if present and not NULL.
 * 
 * @param decompress A flag to specify an additional decompression
 * step if true. */
void DecryptWrapper::initDecryptWrapper(DecryptionKey *deckey, CheckingKey *chkkey, int decompress)
{
   if (deckey)
   {
      switch (deckey->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PRIVATE_RSA:
            { deckeycopy = new RSAPrivateKey(*((RSAPrivateKey *)deckey)); break; }
         case AbstractKey::KEY_TYPE_SYMM_AES:
            { deckeycopy = new AESSymmetricKey(*((AESSymmetricKey *)deckey)); break; }
         default:
            { errprint("decrypt wrapper has no key: can't do decryption!"); deckeycopy = NULL; break; }
      }
   } else
   {
      deckeycopy = NULL;
   }
   if (chkkey)
   {
      switch (chkkey->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PUBLIC_RSA:
            { chkkeycopy = new RSAPublicKey(*((RSAPublicKey *)chkkey)); break; }
         case AbstractKey::KEY_TYPE_PUBLIC_DSA:
            { chkkeycopy = new DSAPublicKey(*((DSAPublicKey *)chkkey)); break; }
         case AbstractKey::KEY_TYPE_PUBLIC_ECDSA:
            { chkkeycopy = new ECDSAPublicKey(*((ECDSAPublicKey *)chkkey)); break; }
         default:
            { errprint("unknown signature check key type"); chkkeycopy = NULL; break; }
      }
   } else
   {
      chkkeycopy = NULL;
   }
   dodecompress = (decompress != 0);
   wastrusted = 0;
   apr_status_t status;
   mp = NULL;
   status = apr_pool_create(&mp, NULL);
   if (mp == NULL)
   {
      errprint("failed to create memory pool");
   }
   CheckAPRError(status);
}

/** Construct a decryption process involving just decryption with the
 * given DecryptionKey.
 * 
 * @param deckey DecryptionKey to use for decrypting data if present
 * and not NULL. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey)
{
   initDecryptWrapper(deckey, NULL, FALSE);
}

/** Construct a decryption process which verifies that the signature
 * on the data matches the given CheckingKey, before decrypting with
 * the given DecryptionKey.
 * 
 * @param deckey DecryptionKey to use for decrypting data if present
 * and not NULL.
 * 
 * @param chkkey CheckingKey to use to verify signatures on the data
 * stream if present and not NULL. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey, CheckingKey *chkkey)
{
   initDecryptWrapper(deckey, chkkey, FALSE);
}

/** Construct a decryption process which decrypts with the given
 * DecryptionKey, but also optionally apply a decompression step to
 * the resulting data.
 * 
 * @param deckey DecryptionKey to use for decrypting data if present
 * and not NULL.
 * 
 * @param dodecompress A flag to specify an additional decompression
 * step if true. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey, int dodecompress)
{
   initDecryptWrapper(deckey, NULL, dodecompress);
}

/** Construct a decryption process which verifies that the signature
 * on the data matches the given CheckingKey, decrypts with the given
 * DecryptionKey and can also also optionally apply a decompression
 * step to the decrypted data.
 * 
 * @param deckey DecryptionKey to use for decrypting data if present
 * and not NULL.
 * 
 * @param chkkey CheckingKey to use to verify signatures on the data
 * stream if present and not NULL.
 * 
 * @param dodecompress A flag to specify an additional decompression
 * step if true. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey, CheckingKey *chkkey, int dodecompress)
{
   initDecryptWrapper(deckey, chkkey, dodecompress);
}

/** Destructor for the decryption wrapper. */
DecryptWrapper::~DecryptWrapper()
{
   if (deckeycopy) delete deckeycopy;
   if (chkkeycopy) delete chkkeycopy;
   apr_pool_clear(mp);
   apr_pool_destroy(mp);
}

/** Apply the decryption process to a buffer of data, returning the
 * plaintext decoded data.
 * 
 * @param buffer Buffer for coded data input.
 * 
 * @param length Length of buffer for encoded data input.
 * 
 * @param[out] newlength Length of plain output buffer
 * 
 * @return A char buffer into which the plain unwrapped data is
 * put. The caller is responsible for deleting the new buffer
 * instance. */
char *DecryptWrapper::decryptDataBuffer(const char *buffer, apr_size_t length, apr_size_t *newlength)
{
   CryptOChain *crypty = new CryptOChain(mp);
   crypty->buildDecryptionChain(deckeycopy, chkkeycopy, dodecompress);
   char *newdata = crypty->processDataBlock(buffer, length, newlength);
   wastrusted = crypty->processWasSecure();
   delete crypty;
   return newdata;
}

/** Apply the decryption process to a buffer of data, returning the
 * plaintext decoded data. Append a magical nul-termination to the
 * buffer, to make it a valid string if need be. This null termination
 * is not counted in the size.
 * 
 * @param buffer Buffer for coded data input.
 * 
 * @param length Length of buffer for encoded data input.
 * 
 * @param[out] newlength Length of plain output buffer
 * 
 * @return A string-terminated char buffer into which the plain
 * unwrapped data is put. The caller is responsible for deleting the
 * new buffer instance. */
char *DecryptWrapper::decryptStringBuffer(const char *buffer, apr_size_t length, apr_size_t *newlength)
{
   char *intermed = decryptDataBuffer(buffer, length, newlength);
   if (intermed == NULL) return NULL;
   char *n_final = (char *)malloc(sizeof(char)*((*newlength) + 1));
   memcpy(n_final, intermed, sizeof(char)*(*newlength));
   n_final[*newlength] = '\0';
   free(intermed);
   return n_final;
}

/** Apply the decryption process to a pipe of data in the form of an
 * APRBucketBrigade, returning a pointer to a newly created
 * APRBucketBrigade containing the decoded data.
 * 
 * @param bbrgde The bucket brigade input of encoded data.
 * 
 * @return New bucket brigade representing the complete decoded data
 * stream. The caller is responsible for deleting the APRBucketBrigade
 * instance. */
APRBucketBrigade *DecryptWrapper::decryptToBucket(APRBucketBrigade *bbrgde)
{
   APRBucketBrigade *outputbrgde = new APRBucketBrigade(mp);
   CryptOChain *crypty = new CryptOChain(*bbrgde, *outputbrgde, mp);
   crypty->buildDecryptionChain(deckeycopy, chkkeycopy, dodecompress);
   crypty->processPipeBegin();
   int prociter = 1;
   // We may need to disallow finishing before a isClosed
   while (prociter)
      prociter = crypty->processPipeStep(ENV_READ_BUFFER, ENV_READ_BUFFER);
   crypty->processPipeEnd();
   wastrusted = crypty->processWasSecure();
   delete crypty;
   if (prociter == -1)
   {
      delete outputbrgde;
      return NULL;
   }
   return outputbrgde;
}

/** Pops the message trusted flag out from the decrypt wrapper and
 * resets the flag.
 * 
 * @return Zero if the message did not get processed for signature or
 * is untrusted, one if the message is trusted. */
int DecryptWrapper::popMessageTrustedFlag()
{
   int lasttrust = wastrusted;
   wastrusted = 0;
   return lasttrust;
}
